#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：T1.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/10 15:01 
'''
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical
import matplotlib.pyplot as plt
import random
from collections import deque

import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict, Optional
from dataclasses import dataclass, field
from enum import Enum
import random
import math
from copy import deepcopy
from abc import ABC, abstractmethod
import sys


class InitializationStrategy(Enum):
    """初始化策略枚举"""
    RANDOM = 1
    HEURISTIC = 2
    PRODUCTIVITY_BASED = 3


@dataclass
class EquipmentParams:
    """设备参数类"""
    failure_rate: float  # 故障率
    repair_rate: float  # 修复率
    capacity_min: float  # 最小产能
    capacity_max: float  # 最大产能
    mttr: float  # 平均修复时间
    mtbf: float  # 平均故障间隔时间


@dataclass
class BufferConfig:
    """缓冲区配置类"""
    line_id: int
    buffer_id: int
    capacity: int
    transfer_rate: float = 0.0
    current_stock: int = 0


@dataclass
class Solution:
    """解决方案类"""
    buffer_configs: List[BufferConfig]
    total_revenue: float = 0.0
    throughput_rate: float = 0.0
    dominance_rank: int = 0
    crowding_distance: float = 0.0
    domination_count: int = 0
    dominated_solutions: List['Solution'] = field(default_factory=list)


class ProductionLineSimulator(ABC):
    """生产线仿真器抽象类"""

    @abstractmethod
    def simulate(self, buffer_configs: List[BufferConfig],
                 simulation_time: float) -> Dict[str, float]:
        """运行仿真"""
        pass

    @abstractmethod
    def calculate_arrival_rate(self, buffer_config: BufferConfig,
                               upstream_equipment: EquipmentParams) -> float:
        """计算到达率"""
        pass

    @abstractmethod
    def calculate_throughput(self, line_id: int,
                             buffer_configs: List[BufferConfig]) -> float:
        """计算吞吐率"""
        pass


class SimplifiedSimulator(ProductionLineSimulator):
    """简化版仿真器（基于数学模型）"""

    def __init__(self, equipment_params: Dict[Tuple[int, int], EquipmentParams],
                 transfer_costs: Dict[Tuple[int, int, int, int], float],
                 product_revenues: Dict[int, float]):
        self.equipment_params = equipment_params
        self.transfer_costs = transfer_costs
        self.product_revenues = product_revenues

    def simulate(self, buffer_configs: List[BufferConfig],
                 simulation_time: float = 30 * 15) -> Dict[str, float]:
        """运行简化仿真"""
        total_revenue = 0.0
        total_throughput = 0.0

        # 按流水线分组
        line_buffers = {}
        for config in buffer_configs:
            if config.line_id not in line_buffers:
                line_buffers[config.line_id] = []
            line_buffers[config.line_id].append(config)

        # 计算各流水线的产出
        for line_id, buffers in line_buffers.items():
            throughput = self.calculate_throughput(line_id, buffers)
            total_throughput += throughput

            # 简化收益计算（假设产品均匀分布）
            avg_revenue = sum(self.product_revenues.values()) / len(self.product_revenues)
            total_revenue += throughput * avg_revenue * simulation_time / 60  # 转换为小时

            # 减去传输成本
            transfer_cost = self._calculate_transfer_cost(buffers)
            total_revenue -= transfer_cost

        return {
            'total_revenue': total_revenue,
            'throughput_rate': total_throughput / len(line_buffers),
            'total_throughput': total_throughput
        }

    def _calculate_transfer_cost(self, buffers: List[BufferConfig]) -> float:
        """计算传输成本"""
        total_cost = 0.0
        for config in buffers:
            if config.transfer_rate > 0:
                # 查找传输成本
                for (src_line, src_buf, dst_line, dst_buf), cost in self.transfer_costs.items():
                    if src_line == config.line_id and src_buf == config.buffer_id:
                        total_cost += cost * config.transfer_rate * 100  # 简化计算
        return total_cost

    def calculate_arrival_rate(self, buffer_config: BufferConfig,
                               upstream_equipment: EquipmentParams) -> float:
        """计算到达率（基于论文公式7）"""
        # 设备可用度
        availability = upstream_equipment.mtbf / (upstream_equipment.mtbf + upstream_equipment.mttr)

        # 缓冲区空缺概率（简化计算）
        p_vacancy = 1 - (buffer_config.current_stock / max(1, buffer_config.capacity))

        # 实际到达率
        arrival_rate = upstream_equipment.capacity_max * availability * p_vacancy

        return arrival_rate

    def calculate_throughput(self, line_id: int,
                             buffer_configs: List[BufferConfig]) -> float:
        """计算吞吐率（基于论文公式11）"""
        # 按工序排序
        buffers_sorted = sorted(buffer_configs, key=lambda x: x.buffer_id)

        throughput = float('inf')
        for config in buffers_sorted:
            # 获取对应设备参数
            equipment_key = (config.line_id, config.buffer_id)
            if equipment_key in self.equipment_params:
                equipment = self.equipment_params[equipment_key]

                # 计算有效产出率（简化版）
                availability = equipment.mtbf / (equipment.mtbf + equipment.mttr)
                buffer_utilization = min(1.0, config.current_stock / max(1, config.capacity))

                line_throughput = equipment.capacity_max * availability * buffer_utilization
                throughput = min(throughput, line_throughput)

        return throughput if throughput != float('inf') else 0.0


class CMOHFA:
    """集群混合多目标萤火虫算法"""

    def __init__(self, num_lines: int, num_buffers_per_line: int,
                 total_buffer_capacity: int, simulator: ProductionLineSimulator):
        self.num_lines = num_lines
        self.num_buffers = num_buffers_per_line
        self.total_capacity = total_buffer_capacity
        self.simulator = simulator

        self.population: List[Solution] = []
        self.archive: List[Solution] = []
        self.buffer_archive: List[Solution] = []

        # 算法参数
        self.pop_size = 60
        self.er_ratio = 0.15
        self.initial_attractiveness = 1.0
        self.attenuation_rate = 1.0
        self.im_parameter = 5
        self.f_check_interval = 5

        # 跟踪变量
        self.iteration_count = 0
        self.inc_num = 0
        self.convergence_history: List[float] = []

    def initialize_population(self):
        """初始化种群（三种混合策略）"""
        self.population = []
        pop_ratios = [0.3, 0.3, 0.4]  # 三种策略的比例

        for strategy, ratio in zip(InitializationStrategy, pop_ratios):
            n_individuals = int(self.pop_size * ratio)

            for _ in range(n_individuals):
                if strategy == InitializationStrategy.RANDOM:
                    solution = self._random_initialization()
                elif strategy == InitializationStrategy.HEURISTIC:
                    solution = self._heuristic_initialization()
                else:
                    solution = self._productivity_based_initialization()

                # 评估解的质量
                self.evaluate_solution(solution)
                self.population.append(solution)

    def _random_initialization(self) -> Solution:
        """随机初始化策略"""
        configs = []
        remaining_capacity = self.total_capacity

        # 为每个缓冲区随机分配容量
        for line in range(self.num_lines):
            for buffer in range(self.num_buffers):
                if remaining_capacity > 0:
                    cap = random.randint(1, min(remaining_capacity, 100))
                    remaining_capacity -= cap
                else:
                    cap = 1

                transfer_rate = random.uniform(0, 0.5)  # 传输率限制在0-0.5
                configs.append(BufferConfig(line, buffer, cap, transfer_rate))

        return Solution(configs)

    def _heuristic_initialization(self) -> Solution:
        """基于启发式的初始化策略"""
        configs = []
        alpha = [0.52, 0.49, 0.47]  # 区域参数
        beta = [0.21, 0.19, 0.23]  # 偏差比例参数

        # 计算权重
        weights = []
        for i in range(self.num_lines):
            for j in range(self.num_buffers):
                mean = alpha[i] * self.num_buffers
                std_dev = beta[i] * mean
                weight = max(0.1, random.gauss(mean, std_dev))
                weights.append(weight)

        # 归一化权重
        total_weight = sum(weights)
        normalized_weights = [w / total_weight for w in weights]

        # 分配容量
        for i, weight in enumerate(normalized_weights):
            line = i // self.num_buffers
            buffer = i % self.num_buffers
            cap = int(weight * self.total_capacity)
            transfer_rate = random.uniform(0, 0.5)
            configs.append(BufferConfig(line, buffer, cap, transfer_rate))

        return Solution(configs)

    def _productivity_based_initialization(self) -> Solution:
        """基于生产率的初始化策略"""
        configs = []

        # 假设的设备产能（实际应从仿真器获取）
        assumed_capacities = []
        for line in range(self.num_lines):
            for buffer in range(self.num_buffers):
                # 简化：随机生成假设产能
                capacity = random.uniform(10, 50)
                assumed_capacities.append(capacity)

        total_capacity = sum(assumed_capacities)

        for i, cap in enumerate(assumed_capacities):
            line = i // self.num_buffers
            buffer = i % self.num_buffers
            alloc_cap = int((cap / total_capacity) * self.total_capacity)
            transfer_rate = random.uniform(0, 0.3)
            configs.append(BufferConfig(line, buffer, alloc_cap, transfer_rate))

        return Solution(configs)

    def evaluate_solution(self, solution: Solution) -> Tuple[float, float]:
        """
        评估解的质量
        返回: (总收益, 吞吐率)
        """
        # 使用仿真器进行评估
        results = self.simulator.simulate(solution.buffer_configs)

        solution.total_revenue = results['total_revenue']
        solution.throughput_rate = results['throughput_rate']

        return solution.total_revenue, solution.throughput_rate

    def non_dominated_sort(self, solutions: List[Solution]) -> List[List[Solution]]:
        """修复后的非支配排序"""
        if not solutions:
            return [[]]

        # 初始化支配关系
        for sol_i in solutions:
            sol_i.domination_count = 0
            sol_i.dominated_solutions = []

        # 计算支配关系
        for i, sol_i in enumerate(solutions):
            for j, sol_j in enumerate(solutions):
                if i != j:
                    if self._dominates(sol_i, sol_j):
                        sol_i.dominated_solutions.append(sol_j)
                    elif self._dominates(sol_j, sol_i):
                        sol_i.domination_count += 1

        # 构建前沿
        fronts = []
        current_front = []

        for sol in solutions:
            if sol.domination_count == 0:
                sol.rank = 0
                current_front.append(sol)

        fronts.append(current_front)

        i = 0
        while fronts[i]:
            next_front = []
            for sol_i in fronts[i]:
                for sol_j in sol_i.dominated_solutions:
                    sol_j.domination_count -= 1
                    if sol_j.domination_count == 0:
                        sol_j.rank = i + 1
                        next_front.append(sol_j)
            i += 1
            fronts.append(next_front)

        # 移除最后一个空的前沿
        if not fronts[-1]:
            fronts.pop()

        return fronts

    def _dominates(self, sol1: Solution, sol2: Solution) -> bool:
        """检查sol1是否支配sol2"""
        # 两个目标都要更好（收益更高且吞吐率更高）
        better_revenue = sol1.total_revenue >= sol2.total_revenue
        better_throughput = sol1.throughput_rate >= sol2.throughput_rate

        # 至少有一个目标严格更好
        strictly_better = (sol1.total_revenue > sol2.total_revenue or
                           sol1.throughput_rate > sol2.throughput_rate)

        return better_revenue and better_throughput and strictly_better

    def _calculate_crowding_distance(self, front: List[Solution]) -> None:
        """计算拥挤度距离"""
        if not front:
            return

        n = len(front)
        for sol in front:
            sol.crowding_distance = 0

        # 对每个目标函数计算拥挤度
        for obj_idx in range(2):  # 两个目标函数
            # 按当前目标排序
            if obj_idx == 0:
                front.sort(key=lambda x: x.total_revenue)
            else:
                front.sort(key=lambda x: x.throughput_rate)

            # 边界解的拥挤度为无穷大
            front[0].crowding_distance = float('inf')
            front[-1].crowding_distance = float('inf')

            # 计算范围
            if len(front) > 2:
                min_val = front[0].total_revenue if obj_idx == 0 else front[0].throughput_rate
                max_val = front[-1].total_revenue if obj_idx == 0 else front[-1].throughput_rate

                if max_val - min_val < 1e-10:  # 避免除零
                    continue

                for i in range(1, n - 1):
                    prev_val = front[i - 1].total_revenue if obj_idx == 0 else front[i - 1].throughput_rate
                    next_val = front[i + 1].total_revenue if obj_idx == 0 else front[i + 1].throughput_rate

                    front[i].crowding_distance += (next_val - prev_val) / (max_val - min_val)

    def firefly_movement(self):
        """萤火虫迁移操作"""
        for i, firefly_i in enumerate(self.population):
            for j, firefly_j in enumerate(self.population):
                if i != j:
                    # 只允许较暗的萤火虫向较亮的移动
                    if self._is_brighter(firefly_j, firefly_i):
                        # 计算吸引度
                        attractiveness = self._calculate_attractiveness(firefly_i, firefly_j)

                        # 执行迁移
                        new_solution = self._move_firefly(firefly_i, firefly_j, attractiveness)

                        # 评估新解
                        self.evaluate_solution(new_solution)

                        # 替换原解（如果新解更好）
                        if self._is_brighter(new_solution, firefly_i):
                            self.population[i] = new_solution

    def _is_brighter(self, sol1: Solution, sol2: Solution) -> bool:
        """检查sol1是否比sol2更亮（更好）"""
        # 基于非支配等级和拥挤度
        if sol1.rank < sol2.rank:
            return True
        elif sol1.rank == sol2.rank and sol1.crowding_distance > sol2.crowding_distance:
            return True
        return False

    def _calculate_attractiveness(self, sol1: Solution, sol2: Solution) -> float:
        """计算吸引度"""
        # 计算欧氏距离（归一化）
        revenues = [s.total_revenue for s in self.population]
        throughputs = [s.throughput_rate for s in self.population]

        revenue_range = max(revenues) - min(revenues) if revenues else 1
        throughput_range = max(throughputs) - min(throughputs) if throughputs else 1

        revenue_dist = abs(sol1.total_revenue - sol2.total_revenue) / max(1, revenue_range)
        throughput_dist = abs(sol1.throughput_rate - sol2.throughput_rate) / max(1, throughput_range)

        total_dist = math.sqrt(revenue_dist ** 2 + throughput_dist ** 2)

        # 吸引度计算
        attractiveness = self.initial_attractiveness * math.exp(
            -self.attenuation_rate * total_dist ** 2)

        return attractiveness

    def _move_firefly(self, source: Solution, target: Solution,
                      attractiveness: float) -> Solution:
        """萤火虫移动"""
        new_configs = []

        for src_cfg, tgt_cfg in zip(source.buffer_configs, target.buffer_configs):
            # 添加高斯扰动
            gaussian_noise = random.gauss(0, 0.1)

            # 计算新容量（整数）
            new_capacity = int(src_cfg.capacity + attractiveness *
                               (tgt_cfg.capacity - src_cfg.capacity) + gaussian_noise)

            # 计算新传输率
            new_transfer_rate = src_cfg.transfer_rate + attractiveness * (
                    tgt_cfg.transfer_rate - src_cfg.transfer_rate) + gaussian_noise * 0.1

            # 确保值在合理范围内
            new_capacity = max(1, min(new_capacity, 100))
            new_transfer_rate = max(0, min(new_transfer_rate, 1))

            new_configs.append(BufferConfig(
                src_cfg.line_id, src_cfg.buffer_id,
                new_capacity, new_transfer_rate
            ))

        return Solution(new_configs)

    def _population_exchange(self):
        """种群信息交换策略"""
        if self.inc_num >= self.im_parameter * 0.5:
            # 合并所有种群
            all_solutions = self.population.copy()
            fronts = self.non_dominated_sort(all_solutions)

            # 基于拥挤度选择最优个体
            new_population = []
            for front in fronts:
                if len(new_population) + len(front) <= self.pop_size:
                    self._calculate_crowding_distance(front)
                    new_population.extend(front)
                else:
                    remaining = self.pop_size - len(new_population)
                    self._calculate_crowding_distance(front)
                    front.sort(key=lambda x: x.crowding_distance, reverse=True)
                    new_population.extend(front[:remaining])
                    break

            self.population = new_population
        else:
            # 随机选择个体进行迁移
            subgroup = []
            for sol in self.population:
                if random.random() < self.er_ratio:
                    subgroup.append(sol)

            # 重新分配子群个体
            for sol in subgroup:
                target_pop_idx = random.randint(0, len(self.population) - 1)
                if self.population[target_pop_idx] != sol:
                    self.population[target_pop_idx] = deepcopy(sol)

    def _update_archive(self, new_front: List[Solution]):
        """更新档案集"""
        # 添加新前沿解
        self.archive.extend(new_front)

        # 重新非支配排序
        fronts = self.non_dominated_sort(self.archive)

        # 只保留最优前沿
        if fronts:
            self.archive = fronts[0]

            # 计算拥挤度
            self._calculate_crowding_distance(self.archive)

            # 如果档案过大，基于拥挤度修剪
            if len(self.archive) > 2 * self.pop_size:
                self.archive.sort(key=lambda x: x.crowding_distance, reverse=True)
                self.archive = self.archive[:2 * self.pop_size]

    def _check_convergence(self) -> bool:
        """检查收敛条件"""
        if self.inc_num >= self.im_parameter:
            return True

        # 检查前沿是否更新
        if len(self.convergence_history) >= 2:
            improvement = abs(self.convergence_history[-1] - self.convergence_history[-2])
            if improvement < 1e-5:  # 收敛阈值
                self.inc_num += 1
            else:
                self.inc_num = 0
        else:
            self.inc_num = 0

        return False

    def run_optimization(self, max_iterations: int = 50) -> Dict:
        """运行优化过程"""
        print("初始化种群...")
        self.initialize_population()

        print("开始优化...")
        for iteration in range(max_iterations):
            self.iteration_count = iteration

            # 非支配排序
            fronts = self.non_dominated_sort(self.population)

            if fronts and fronts[0]:
                # 更新档案集
                self._update_archive(fronts[0])

                # 记录收敛历史
                best_revenue = max(sol.total_revenue for sol in fronts[0])
                self.convergence_history.append(best_revenue)

            # 萤火虫迁移
            self.firefly_movement()

            # 种群信息交换
            self._population_exchange()

            # 检查收敛条件
            if self._check_convergence():
                print(f"算法在迭代 {iteration} 次后收敛")
                break

            if iteration % 10 == 0:
                best_front = fronts[0] if fronts and fronts[0] else []
                best_revenue = max(sol.total_revenue for sol in best_front) if best_front else 0
                best_throughput = max(sol.throughput_rate for sol in best_front) if best_front else 0
                print(f"迭代 {iteration}: 最优收益={best_revenue:.2f}, 最优吞吐率={best_throughput:.4f}")

        return {
            'archive': self.archive,
            'convergence_history': self.convergence_history,
            'final_population': self.population
        }


def create_experiment_setup():
    """创建实验配置"""
    # 设备参数（示例数据）
    equipment_params = {}
    for line in range(3):  # 3条流水线
        for buffer in range(8):  # 每条线8个缓冲区
            equipment_params[(line, buffer)] = EquipmentParams(
                failure_rate=random.uniform(0.01, 0.05),
                repair_rate=random.uniform(0.1, 0.3),
                capacity_min=random.uniform(5, 10),
                capacity_max=random.uniform(20, 30),
                mttr=random.uniform(10, 30),
                mtbf=random.uniform(100, 200)
            )

    # 传输成本（示例数据）
    transfer_costs = {}
    for src_line in range(3):
        for src_buf in range(8):
            for dst_line in range(3):
                if src_line != dst_line:  # 只在不同流水线间传输
                    for dst_buf in range(8):
                        cost = random.uniform(5, 20)
                        transfer_costs[(src_line, src_buf, dst_line, dst_buf)] = cost

    # 产品收益（示例数据）
    product_revenues = {0: 100, 1: 150, 2: 200, 3: 120, 4: 180}

    return equipment_params, transfer_costs, product_revenues


def run_complete_experiment():
    """运行完整实验"""
    print("创建实验配置...")
    equipment_params, transfer_costs, product_revenues = create_experiment_setup()

    print("初始化仿真器...")
    simulator = SimplifiedSimulator(equipment_params, transfer_costs, product_revenues)

    print("创建优化算法实例...")
    algorithm = CMOHFA(
        num_lines=3,
        num_buffers_per_line=8,
        total_buffer_capacity=1500,
        simulator=simulator
    )

    print("运行优化算法...")
    results = algorithm.run_optimization(max_iterations=50)

    return results, algorithm


def analyze_results(results: Dict, algorithm: CMOHFA):
    """分析实验结果"""
    # 提取帕累托前沿
    pareto_front = results['archive']

    if not pareto_front:
        print("未找到帕累托解")
        return

    # 绘制帕累托前沿
    revenues = [sol.total_revenue for sol in pareto_front]
    throughputs = [sol.throughput_rate for sol in pareto_front]

    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.scatter(throughputs, revenues, alpha=0.7, c='blue')
    plt.xlabel('吞吐率')
    plt.ylabel('总收益')
    plt.title('帕累托前沿')
    plt.grid(True)

    # 绘制收敛曲线
    plt.subplot(1, 2, 2)
    plt.plot(results['convergence_history'])
    plt.xlabel('迭代次数')
    plt.ylabel('最佳收益')
    plt.title('算法收敛曲线')
    plt.grid(True)

    plt.tight_layout()
    plt.show()

    # 输出最优解信息
    best_solution = max(pareto_front, key=lambda x: x.total_revenue)
    print(f"\n最优解结果:")
    print(f"总收益: {best_solution.total_revenue:.2f}")
    print(f"吞吐率: {best_solution.throughput_rate:.4f}")

    # 输出缓冲区配置
    print("\n缓冲区配置详情:")
    for config in best_solution.buffer_configs:
        print(f"流水线{config.line_id}-缓冲区{config.buffer_id}: "
              f"容量={config.capacity}, 传输率={config.transfer_rate:.3f}")


# 测试用例
def test_algorithm():
    """测试算法基本功能"""
    print("=== 测试算法基本功能 ===")

    # 创建简单的测试配置
    equipment_params = {
        (0, 0): EquipmentParams(0.01, 0.2, 5, 25, 15, 150),
        (0, 1): EquipmentParams(0.02, 0.25, 6, 28, 12, 140)
    }

    transfer_costs = {
        (0, 0, 1, 0): 10.0,
        (0, 1, 1, 0): 12.0
    }

    product_revenues = {0: 100, 1: 150}

    simulator = SimplifiedSimulator(equipment_params, transfer_costs, product_revenues)

    # 创建小规模算法实例
    algorithm = CMOHFA(
        num_lines=1,
        num_buffers_per_line=2,
        total_buffer_capacity=100,
        simulator=simulator
    )

    # 测试初始化
    algorithm.initialize_population()
    print(f"种群大小: {len(algorithm.population)}")

    # 测试评估
    test_solution = algorithm.population[0]
    revenue, throughput = algorithm.evaluate_solution(test_solution)
    print(f"测试解收益: {revenue:.2f}, 吞吐率: {throughput:.4f}")

    # 测试非支配排序
    fronts = algorithm.non_dominated_sort(algorithm.population[:10])  # 只测试前10个解
    print(f"前沿数量: {len(fronts)}")

    print("=== 测试完成 ===")


# 运行主程序
if __name__ == "__main__":
    # 运行测试
    test_algorithm()

    print("\n" + "=" * 50)
    print("开始完整实验...")
    print("=" * 50)

    # 运行完整实验
    results, algorithm = run_complete_experiment()

    # 分析结果
    analyze_results(results, algorithm)
